import re
from typing import List, Optional, Union, BinaryIO
from dataclasses import dataclass
from pypdf import PdfReader
import io

@dataclass
class SplitConfig:
    """文本分割配置"""
    strategy: str = "paragraph"  # paragraph, length, sentence, markdown, sliding
    chunk_size: int = 1000      # 分块大小(字符数)
    chunk_overlap: int = 200    # 滑动窗口重叠大小
    separator: str = "\n\n"     # 分隔符

class TextSplitter:
    def __init__(self, config: SplitConfig):
        self.config = config

    def process_file(self, file_content: Union[str, bytes, BinaryIO], file_type: str) -> List[str]:
        """处理不同类型的文件内容并返回分片结果"""
        if file_type == 'pdf':
            text = self._extract_text_from_pdf(file_content)
        else:
            # 假设是纯文本
            text = file_content if isinstance(file_content, str) else file_content.decode('utf-8')
        
        return self.split(text)

    def _extract_text_from_pdf(self, file_content: Union[bytes, BinaryIO]) -> str:
        """从PDF文件中提取文本"""
        if isinstance(file_content, bytes):
            file_content = io.BytesIO(file_content)
            
        pdf = PdfReader(file_content)
        text_parts = []
        
        for page in pdf.pages:
            text = page.extract_text()
            if text:
                text_parts.append(text)
        
        return "\n\n".join(text_parts)

    def split(self, text: str) -> List[str]:
        if self.config.strategy == "paragraph":
            return self._split_by_paragraph(text)
        elif self.config.strategy == "length":
            return self._split_by_length(text)
        elif self.config.strategy == "sentence":
            return self._split_by_sentence(text)
        elif self.config.strategy == "markdown":
            return self._split_by_markdown(text)
        elif self.config.strategy == "sliding":
            return self._split_sliding_window(text)
        else:
            raise ValueError(f"Unknown split strategy: {self.config.strategy}")

    def _split_by_paragraph(self, text: str) -> List[str]:
        # 将文本中的 \r\n 统一转换为 \n
        text = text.replace('\r\n', '\n')
        # 使用两个或更多换行符进行分割
        paragraphs = [p.strip() for p in re.split(r'\n\s*\n', text) if p.strip()]
        return paragraphs if paragraphs else [text]

    def _split_by_length(self, text: str) -> List[str]:
        chunks = []
        current_chunk = []
        current_length = 0
        
        for char in text:
            current_chunk.append(char)
            current_length += 1
            
            if current_length >= self.config.chunk_size:
                chunks.append("".join(current_chunk))
                current_chunk = []
                current_length = 0
                
        if current_chunk:
            chunks.append("".join(current_chunk))
        return chunks

    def _split_by_sentence(self, text: str) -> List[str]:
        # 简单的句子分割规则，可以根据需要优化
        sentence_endings = r'[.!?。！？]+'
        sentences = re.split(f'({sentence_endings})', text)
        result = []
        current = []
        
        for i in range(0, len(sentences)-1, 2):
            if sentences[i]:
                current.append(sentences[i] + (sentences[i+1] if i+1 < len(sentences) else ''))
                if sum(len(s) for s in current) >= self.config.chunk_size:
                    result.append(''.join(current))
                    current = []
        
        if current:
            result.append(''.join(current))
        return result

    def _split_by_markdown(self, text: str) -> List[str]:
        # 按markdown标题分割
        header_pattern = r'^#{1,6}\s+.*$'
        sections = re.split(f'({header_pattern})', text, flags=re.MULTILINE)
        chunks = []
        current_section = []
        
        for section in sections:
            if section.strip():
                current_section.append(section)
                if len(''.join(current_section)) >= self.config.chunk_size:
                    chunks.append(''.join(current_section))
                    current_section = []
                    
        if current_section:
            chunks.append(''.join(current_section))
        return chunks

    def _split_sliding_window(self, text: str) -> List[str]:
        # 滑动窗口分割，有重叠部分
        chunks = []
        start = 0
        while start < len(text):
            end = start + self.config.chunk_size
            chunk = text[start:end]
            chunks.append(chunk)
            start += self.config.chunk_size - self.config.chunk_overlap
        return chunks